# -*- coding: utf-8 -*-
"""
Created on Sun Jan  3 01:19:55 2016

@author: liuyi05
"""

import logging
import pandas as pd
from core.account import Account, Delegation, Op
from core.calc import Calc
from core.strategy import Turtle, MovingAverage

def simulate(code, df, account):
    assert("op" in df)
    atr20d = "atr20d"
    print("code:", code)
    if Calc.LABEL_RETURN not in df:
        Calc.calc_return(df)
    if atr20d not in df:
        Calc.calc_atrxd(df, 20, atr20d)
    theta = 0.09
    eps = 1e-6
    ret = df[Calc.LABEL_RETURN]
    allot = df[Calc.LABEL_ALLOT]
    atr = df[atr20d]
    open_price = df["open"]
    high = df["high"]
    close = df["close"]
    low = df["low"]
    op = df["op"]
    date = df.index
    market_price = {}
    blc = pd.Series(0, index=df.index)
    mv = pd.Series(0, index=df.index)
    mp = pd.Series(0, index=df.index)
    ass = pd.Series(0, index=df.index)
    fee = pd.Series(0, index=df.index)
    qnt = pd.Series(0, index=df.index)
    for i in range(df.shape[0]):
        if allot[i] > eps:
            account.allot_shares(code, allot[i])
            print("allotment")

        operation = op[i]
        if account.contains(code) and open_price[i] < \
                account.get_share(code).loss_limit(atr[i]):
            operation = -1
            print("limit loss")
        
        if operation > eps:
            quant = account.bal_rate2quant(open_price[i], operation)
            if quant > 0:
                if high[i] == low[i] and ret[i] > theta:
                    print("limit up")
                else:
                    delegation = Delegation(Op.BUY, date[i], code,
                                            open_price[i], quant, True)
                    trans = account.delegate(delegation)
                    if trans is not None:
                        print("Trans:", trans)

        if operation < -eps:
            if account.contains(code):
                if high[i] == low[i] and ret[i] < -theta:
                    print("limit down")
                else:
                    quant = account.pos_rate2quant(code, -operation)
                    delegation = Delegation(Op.SELL, date[i], code,
                                            open_price[i], quant, True)
                    trans = account.delegate(delegation)
                    if trans is not None:
                        print("Trans:", trans)

        blc[i] = account.balance
        market_price[code] = close[i]
        mv[i] = account.market_value(market_price)
        mp[i] = account.market_profit(market_price)
        ass[i] = account.balance + mv[i]
        fee[i] = account.accum_fee
        qnt[i] = account.get_quant(code)

    return pd.DataFrame({"balance": blc, "market_value": mv,
                         "market_profit": mp, "asset": ass, "cum_fee": fee,
                         "position": qnt})

def main():
    code = "600519"
    #code = "000078"
    dir_path = "data/history/040101_051130/"
    path = dir_path + code + ".asc"
    df = pd.read_csv(path, index_col="date")

    strategy = Turtle(10, 20)
    # TODO: 第一次购买信号不置信
    #strategy = MovingAverage(20, 60)
    #strategy = KDJ(10)
    strategy.gen_op_series(df)
    df.dropna(how="any", inplace=True)
    df.to_csv("haha")
    start = 1000000
    account = Account(start)
    state = simulate(code, df, account)
    end = state["asset"][-1]
    df = pd.concat([df, state], axis=1)
    df.to_csv("haha.csv")
    df["asset"].plot(style='r')
    print("SUMMARY:\nstart=[%.2f] end=[%.2f] profit=[%.2f] ratio=[%.2f]" %
            (start, end, end - start, (end - start) / start))

if __name__ == "__main__":
    logging.basicConfig(filename="stock.log")
    logging.getLogger().addHandler(logging.StreamHandler())
    main()